# import matplotlib
# matplotlib.use('Agg')
import copy
import datetime
import os

from torchinfo import summary

from utils.options import args_parser
from models.Update import *
# from models.resnet import *
from models.ResNet_Dery import *
from models.model_creator import *
from models.test import *
from models.aggregation import *
from models.branchnet import *
from models.cnn import *
from models.MobileNetV3 import *
from algorithm.Training_FedAvg import *

from algorithm.Training_FedDF import *

from utils.get_dataset import get_dataset
from utils.save_result import save_result
from utils.set_seed import set_random_seed
from sim.block_meta import MODEL_BLOCKS,MODEL_ZOO,PLANES
from collections import OrderedDict
# args = args_parser()
# net_zoo = []
# net1 = ResNet18_cifar10(num_classes = args.num_classes)
# net2 = ResNet50_cifar10(num_classes = args.num_classes)
# net3 = MobileNetV3(num_classes = args.num_classes)
# net_zoo.append(net1)
# net_zoo.append(net2)
# net_zoo.append(net3)
# # print(net1)
# # print(type(net1))
# def split_block(net_zoo,num_block):
#     block_list = []  ## size = num_block
#     for i in range(len(net_zoo)):
#         block_list.append([])
#         for j in range(num_block):
#             block_list[i].append([])
#             for layer in MODEL_BLOCKS[MODEL_ZOO[i]][j]:
#                 if '.' not in layer:
#                     layer1 = layer
#                     layer2 = None
#                 else:
#                     layers = layer.split('.')
#                     layer1 = layers[0]
#                     layer2 = layers[1]
#                 for name,child in net_zoo[i].named_children():
#                     if layer1 in name :
#                         if isinstance(child,torch.nn.Sequential) :
#                             print(int(layer2))
#                             block_list[i][j].append(child[int(layer2)])
#                         else:
#                             block_list[i][j].append(child)
#     for i in range(len(net_zoo)):
#         for j in range(num_block):
#             block_list[i][j] = nn.Sequential(*block_list[i][j])
#     return block_list
# def split_block_simple(net_zoo,num_block):
#     block_list = []
#     for i in range(len(net_zoo)):
#         block_list.append([])
#         for j in range(num_block):
#             block_list[i].append([])
#             for name,child in net_zoo[i].named_children():
#                 if 'block{:d}'.format(j+1) in name:
#                     block_list[i][j].append(child)
#     return block_list  

# # for name, child in net1.named_children():
# #     print(name)
# #     if isinstance(child, torch.nn.Sequential):
# #                   print(name,child[0])
# #             # for sub_name, sub_child in child.named_children():
# #                 #   print(sub_name,sub_child,type(sub_child),sep=" ")
# #     else:
# #         print(child, type(child), sep=" ")


# # block_list = split_block(net_zoo,4)
# # temp_list = []
# # for i in range(4):
# #     temp_list.append(block_list[2][i])
# # choose = [2,2,2,2]
# # mymodel = CombinedSingleModel(temp_list)
# # print(mymodel)


# # net_test = []
# # net_test.append(mymodel)
# # net_test.append(mymodel)
# # block_list_s = split_block_simple(net_test,num_block=4)
# # print(block_list_s[0])



# # print(type(mymodel.state_dict()))
# # print(mymodel.state_dict()['block1.0.weight'])
# # # for name,para in mymodel.named_parameters():
# # #     print(name,para,sep=" ")
# # l = [OrderedDict()]
# # for key,name in mymodel.state_dict().items():
# #     if 'block1' in key:
# #         l[0][key] = name
# # print(l[0],sep=" ")

# # net1 = ResNet18_cifar10(num_classes = args.num_classes)
# # net1 = MobileNetV3(num_classes = args.num_classes)
# # for name,child in net1.named_children():
# #     print(name,sep=" ")
# # from models.ViT_T import *
# # net1 = ViT_T12_cifar10(num_classes = 10)
# # print(net1)
# import torchvision
# from torchinfo import summary
from models.MobileNetV2 import *
# from models.Resnet20 import *
# from models.ResNet8 import *
# # net = torchvision.models.regnet_y_800mf()
# from models.Vgg11 import *# 注意：模型内部传参数和不传参数，输出的结果是不一样的
# # 计算网络参数
# # net = ResNet101_cifar(args)
# # net = MobileNetV2(num_classes = 10)
# # net = MobileNetV2(num_classes = 10)
# # net = ResNet20_cifar(args)
# # net = ResNet8(num_classes = 10)
# # input = torch.rand(50, 3, 32, 32)
# # summary(net, (50, 3, 32, 32))
# # net = MobileNetV3(num_classes = 100)
# # total = sum([param.nelement() for param in net.parameters()])
# # # 精确地计算：1MB=1024KB=1048576字节
# # print('Number of parameter: % .4fM' % (total / 1e6))
args = args_parser()
# net = mobilenetv2(args,alpha=1)
net = MobileNetV3()
input = torch.rand(50, 3, 32, 32)
# summary(net, (50, 3, 32, 32))

data = (net,1)
a= data[0]
b = data[1]
print(a)
print(b)